#!/usr/bin/python

"""
Author:			Jimmy Saw
Date modified:	07-08-2012

Description:    This program draws recruitment plots from BLASTn coordinates.

Usage:          dissertation_RecruitmentPlot.py <BLASTn output file> <genome fasta>

Example:        dissertation_RecruitmentPlot.py genome.blastn genome.fasta

Note: 			BLASTn needs to be run with -m 8 option
"""

import sys
import matplotlib.pyplot as plt
import pylab
import numpy as np
from Bio import SeqIO
from Bio.SeqUtils import GC

def calAvg(countlist):
    total = 0
    count = len(countlist)
    for i in countlist:
        total += i
    avg = float(total/count)
    return avg

filename = sys.argv[1]
f = open(filename, "rU")
fl = f.readlines()

seqfile = sys.argv[2]
sf = SeqIO.read(seqfile, "fasta")
gc_values = []
gcx = []
genome_size = len(sf.seq)

i = 0

while i < len(sf.seq):
    gc = GC(sf.seq[i:i+1000])
    gc_values.append(gc)
    gcx.append(i)
    i += 1000

x1 = []
x2 = []
y = []
col = []

genome_coords = {}

for i in range(1, genome_size+1):
    genome_coords[i] = 0

toplot = []

for i, line in enumerate(fl):
    l = line.split('\t')
    start = int(l[6])
    stop = int(l[7])
    ident = float(l[2])
    x1.append(start)
    x2.append(stop)
    y.append(ident)
    fillcolor = '#AAAAAA'
    if ident >= 90:
        fillcolor = '#FF0000'
    elif ident >= 80:
        fillcolor = '#71C671'
    elif ident >= 70:
        fillcolor = '#7171C6'
    elif ident >= 60:
        fillcolor = '#CDB5CD'
    else:
        fillcolor = '#C5C1AA'
    col.append(fillcolor)
    for a in range(start, stop):
        genome_coords[a] = genome_coords[a] + 1
    toplot.append((start, stop, ident, fillcolor))

gx = [1,genome_size]
gy = [10,10]

xcoords = range(1, genome_size+1)

#sequence coverage
cov1 = []

for k, v in genome_coords.iteritems():
    cov1.append(v)

w = 1000

depthx1 = []
depthy1 = []
while w < len(cov1):
    average = calAvg(cov1[w-1000:w])
    depthx1.append(w-1000+500)
    depthy1.append(average)
    w += 1000 #iterations need to be >1000 to prevent computer crash

print "Plotting..."

fig = plt.figure(1, figsize=(12,4))
for x in toplot:
    plt.plot([x[0], x[1]], [x[2], x[2]], color=x[3], linestyle='-', linewidth=1.5)
plt.axis([0, genome_size, 0, 110])
#plt.title('Recruitment plot')
plt.ylabel('% Identities')
plt.grid(True)

ax2 = plt.twinx()
#ax2.plot(gcx, gc_values, color='green', linestyle='-', alpha=0.4)
ax2.fill_between(depthx1, depthy1, facecolor='#336600', edgecolor='#336600', alpha=0.5)
ax2.axis([0, genome_size, 0, 110])
ax2.set_ylabel('Avg reads recruited (per 1kb)')

plt.show()

